#include "NRreger.h"


NRreger::NRreger(/* args */)
{
}

NRreger::~NRreger()
{
    delete dataloder_;
    delete dem_main_;
    delete dem_tar_;
}

NRreger::NRreger(DataLoder* data,Paras paras)
{
    dataloder_ = data;
    pars_ = paras;
    
    dem_main_ = new DEM();  
    std::shared_ptr<open3d::geometry::PointCloud> pc (new open3d::geometry::PointCloud);
    *pc = dataloder_->PcRecover(0);
    dem_main_->DEM_generate(pc,pars_.DEM_resolusion);
    
    std::shared_ptr<open3d::geometry::PointCloud> ptrpc (new open3d::geometry::PointCloud); 
    *ptrpc = dem_main_->Generate_pc();;
    open3d::visualization::DrawGeometries({ptrpc},"PC",640,480,50,50);
    
}

void NRreger::Initial_GPU(int t)
{
    
    cimg_tar_ = dataloder_->Getframes(t+pars_.keyframe_inte).c_img;
    dimg_tar_ = dataloder_->Getframes(t+pars_.keyframe_inte).d_img;
    cimg_src_ = dataloder_->Getframes(t).c_img;
    dimg_src_ = dataloder_->Getframes(t).d_img;

    dem_tar_ = new DEM();
    std::shared_ptr<open3d::geometry::PointCloud> pc_tar (new open3d::geometry::PointCloud);
    *pc_tar = dataloder_->PcRecover(t+pars_.keyframe_inte);
    dem_tar_->DEM_generate(pc_tar,pars_.DEM_resolusion);
    pc_DEM_tar_ = dem_tar_->Generate_pc();

    dem_main_->ExtractWarppc(2);
    std::string flowpath = str_format("/home/mx/桌面/data/flows-skip7/%05d.txt",t+pars_.n_strat_+pars_.keyframe_inte);
    Eigen::Vector3d initialv = FlowCorrespondence(correspondence_pairs_,flowpath);
    
    //矩阵size
    n_warp_num_ = dem_main_->Getwarppc().points_.size();//m
    n_corr_num_ = correspondence_pairs_.size();//n
    n_allpoint_num = dem_main_->GetMap_e().size();//a
    //申请内存
    f_ = (double*)calloc(n_corr_num_*n_warp_num_,sizeof(double));
    x_ = (double*)calloc(n_warp_num_*3,sizeof(double));
    s_t_ = (double*)calloc(n_corr_num_*3,sizeof(double));
    e_w_ = (double*)calloc(n_warp_num_*n_warp_num_,sizeof(double));
    m_ = (double*)calloc(n_allpoint_num*n_warp_num_,sizeof(double));
    ap_tps_ = (double*)calloc(n_allpoint_num*3,sizeof(double));//ap_tps要跟随tps变化
    grad_ = (double*)calloc(n_warp_num_*3,sizeof(double));

    cudaStat_ = cudaMalloc((void**)&f_d,n_corr_num_*n_warp_num_*sizeof(*f_));
    cudaStat_ = cudaMalloc((void**)&x_d,n_warp_num_*3*sizeof(*x_));
    cudaStat_ = cudaMalloc((void**)&s_t_d,n_corr_num_*3*sizeof(*s_t_));
    cudaStat_ = cudaMalloc((void**)&e_w_d,n_warp_num_*n_warp_num_*sizeof(*e_w_));
    cudaStat_ = cudaMalloc((void**)&m_d,n_allpoint_num*n_warp_num_*sizeof(*m_));
    cudaStat_ = cudaMalloc((void**)&ap_tps_d,n_allpoint_num*3*sizeof(*ap_tps_));
    //赋值
    // f,ew,m的赋值 以及x的初始化
    dem_main_->Coffset_GPU(f_,n_corr_num_,e_w_,n_warp_num_,m_,n_allpoint_num,correspondence_pairs_,pars_.searchratio,x_);
    //x 的初始化
// #pragma omp parallel for
//     for(int i=0;i<n_warp_num_;i++)
//     {
//         x_[i]=initialv(0);
//         x_[i+n_warp_num_]=initialv(1);
//         x_[i+n_warp_num_*2]=initialv(2);
//     }
    //s_t_
#pragma omp parallel for
    for(int i=0;i<n_corr_num_;i++)
    {
        s_t_[i]=correspondence_pairs_[i].position_src(0)-correspondence_pairs_[i].position_tar(0);
        s_t_[i+n_corr_num_]=correspondence_pairs_[i].position_src(1)-correspondence_pairs_[i].position_tar(1);
        s_t_[i+2*n_corr_num_]=correspondence_pairs_[i].position_src(2)-correspondence_pairs_[i].position_tar(2);
    }
    
    //初始化所有点的位置
    deformap = (double*)malloc(n_allpoint_num*3*sizeof(double));
    initialap = (double*)malloc(n_allpoint_num*3*sizeof(double));
    auto pc_all = dem_main_->Generate_pc();
#pragma omp parallel for    
    for(size_t i=0;i<n_allpoint_num;i++)
    {
        deformap[i]=pc_all.points_[i](0);
        deformap[n_allpoint_num+i]=pc_all.points_[i](1);
        deformap[n_allpoint_num*2+i]=pc_all.points_[i](2);
        initialap[i]=pc_all.points_[i](0);
        initialap[n_allpoint_num+i]=pc_all.points_[i](1);
        initialap[n_allpoint_num*2+i]=pc_all.points_[i](2);
    }
    DEM* dem_tps = new DEM();
    std::shared_ptr<open3d::geometry::PointCloud> pc_tar_dem (new open3d::geometry::PointCloud);
    *pc_tar_dem = pc_DEM_tar_;
    dem_tps->DEM_generate(pc_tar_dem,pars_.TPS_resolusion);
    pc_tps = dem_tps->Generate_pc();
    ntps_ = pc_tps.points_.size();
    TPS_GPU(pc_tps);//算出tps参数，表示目标平面
    delete dem_tps;

    //为所有参数在cuda空间赋值
    stat_ = cublasCreate(&handle_);
    stat_ = cublasSetMatrix(n_corr_num_,n_warp_num_,sizeof(double),f_,n_corr_num_,f_d,n_corr_num_);
    stat_ = cublasSetMatrix(n_warp_num_,n_warp_num_,sizeof(double),e_w_,n_warp_num_,e_w_d,n_warp_num_);
    stat_ = cublasSetMatrix(n_allpoint_num,n_warp_num_,sizeof(double),m_,n_allpoint_num,m_d,n_allpoint_num);
}

void NRreger::Clear()
{
    correspondence_pairs_.clear();
    cudaFree ( f_d );
    cudaFree ( x_d );
    cudaFree ( s_t_d );
    cudaFree ( e_w_d );
    cudaFree ( m_d );
    cudaFree ( ap_tps_d );
    cublasDestroy ( handle_ );
    free ( f_ );
    free ( x_ );
    free ( s_t_ );
    free ( e_w_ );
    free ( m_ );
    free ( ap_tps_ );
    free(X_);
    free(deformap);
    free(initialap);
    free(grad_);

}

Eigen::Vector3d NRreger::FlowCorrespondence(VPairs & corres,std::string flowpath){
    Eigen::Vector3d v3d;
    v3d.setZero();

    std::ifstream infile;//文件流
    std::istringstream iss;
    std::string s;
    int delta;
    std::vector<cv::Point> a,b;

    infile.open(flowpath);
    int uv_cout = 0;
    while(getline(infile, s)) 
    {   
        
        Closest c;
        int u =floor(uv_cout/448.0);//行数
        int v = uv_cout-448*u;
        c.src_uv.x = v;
        c.src_uv.y = u;
        iss.clear();
        iss.str(s);
        int count=0;
        // 逐词读取，遍历每一行中的每个词 t
        while(iss>>delta) 
        {
            if(count==0)
            {
                c.tar_uv.x = v+delta;//这里是正向加，反向减
            }
            else{
                c.tar_uv.y = u+delta;
            }
            count++;

        }

        Eigen::Vector3d pix_coor_src(c.src_uv.x,c.src_uv.y,1);
        Eigen::Vector3d pix_coor_tar(c.tar_uv.x,c.tar_uv.y,1);
        uchar z_src = 0; 
        uchar z_tar = 0;
        if(c.src_uv.x<256&&c.tar_uv.x<256&&c.src_uv.y<448&&c.tar_uv.y<448 && c.src_uv.x>=0&&c.tar_uv.x>=0&&c.src_uv.y>=0&&c.tar_uv.y>=0)
        {
            z_src = dimg_src_.at<uchar>(c.src_uv);
            z_tar = dimg_tar_.at<uchar>(c.tar_uv);
        }

        if(z_src!=0&&z_tar!=0)
        {
            c.position_tar = z_tar*pars_.K.inverse()*pix_coor_tar;
            c.position_src = z_src*pars_.K.inverse()*pix_coor_src;
                      
            if(uv_cout%pars_.flow_filter==0)
            {
                corres.push_back(c);
                Eigen::Vector3d deltav = c.position_tar - c.position_src;  
                v3d += deltav;
            }
            
            a.push_back(c.src_uv); 
            b.push_back(c.tar_uv);
        }
        uv_cout++;
        
    }
    v3d = v3d/corres.size();
    // cor_line(a,b,cimg_src_,cimg_tar_);
    return v3d;

    
}

int NRreger::DoNonRigid()
{
    Scalar data_err, smooth_err,tps_err;
    int total_inner_iters = 0; 
    total_inner_iters += DecentSolver_GPU(data_err, smooth_err, tps_err);
    return total_inner_iters;
                                              
}

Scalar NRreger::sample_gradient_GPU(Scalar &data_err, Scalar &smooth_err,Scalar &tps_err,bool calgrad)
{
    //每次循环std矩阵的值需要重新赋值，x_d也需要重新赋值
    stat_ = cublasSetMatrix(n_corr_num_,3,sizeof(double),s_t_,n_corr_num_,s_t_d,n_corr_num_);
    stat_ = cublasSetMatrix(n_warp_num_,3,sizeof(double),x_,n_warp_num_,x_d,n_warp_num_);

    //梯度
    double *r1,*r1_d,*r2,*r2_d,*r3,*r3_d;
    r1 = (double*)calloc(n_warp_num_*3,sizeof(double));
    r2 = (double*)calloc(n_warp_num_*3,sizeof(double));
    r3 = (double*)calloc(n_warp_num_*3,sizeof(double));
    cudaStat_ = cudaMalloc((void**)&r1_d,n_warp_num_*3*sizeof(*r1));
    stat_ = cublasSetMatrix(n_warp_num_,3,sizeof(*r1),r1,n_warp_num_,r1_d,n_warp_num_);
    cudaStat_ = cudaMalloc((void**)&r2_d,n_warp_num_*3*sizeof(*r2));
    stat_ = cublasSetMatrix(n_warp_num_,3,sizeof(*r2),r2,n_warp_num_,r2_d,n_warp_num_);
    cudaStat_ = cudaMalloc((void**)&r3_d,n_warp_num_*3*sizeof(*r3));
    stat_ = cublasSetMatrix(n_warp_num_,3,sizeof(*r3),r3,n_warp_num_,r3_d,n_warp_num_);

    double *data_error,*smooth_error,*tps_error;//能量向量
    data_error = (double*)calloc(n_corr_num_*3,sizeof(double));
    smooth_error = (double*)calloc(n_warp_num_*3,sizeof(double));
    tps_error = (double*)calloc(n_allpoint_num*3,sizeof(double));
    
    double al = 2*pars_.alpha;
    double be = 2*pars_.beta;
    double ga = 2*pars_.tps;
    double beta = 1.0;
    double gama = 0.0;
    //计算 FX+S-T
    stat_ = cublasDgemm(handle_,CUBLAS_OP_N,CUBLAS_OP_N,n_corr_num_,3,n_warp_num_,&beta,f_d,n_corr_num_,x_d,n_warp_num_,&beta,s_t_d,n_corr_num_);
    stat_ = cublasGetMatrix(n_corr_num_ ,3, sizeof (double) ,s_t_d ,n_corr_num_ ,data_error , n_corr_num_ );    /*传出计算能量大小*/

    //计算 F^T*(FX+S-T)，计算结果在s_t_d;
    if(calgrad)
    {
        stat_ = cublasDgemm(handle_,CUBLAS_OP_T,CUBLAS_OP_N,n_warp_num_,3,n_corr_num_,&al,f_d,n_corr_num_,s_t_d,n_corr_num_,&gama,r1_d,n_warp_num_);
        stat_ = cublasGetMatrix(n_warp_num_ ,3, sizeof (* r1_d ) ,r1_d ,n_warp_num_ ,r1 , n_warp_num_ );//得到梯度1
    }

    //计算 (E-W)^T*(E-W)*X

    stat_ = cublasDgemm(handle_,CUBLAS_OP_N,CUBLAS_OP_N,n_warp_num_,3,n_warp_num_,&beta,e_w_d,n_warp_num_,x_d,n_warp_num_,&gama,r2_d,n_warp_num_);
    stat_ = cublasGetMatrix(n_warp_num_ ,3, sizeof(double) ,r2_d ,n_warp_num_ ,smooth_error , n_warp_num_ );//平滑能量

    if(calgrad)
    {
        stat_ = cublasDgemm(handle_,CUBLAS_OP_T,CUBLAS_OP_N,n_warp_num_,3,n_warp_num_,&be,e_w_d,n_warp_num_,r2_d,n_warp_num_,&gama,r2_d,n_warp_num_);
        stat_ = cublasGetMatrix(n_warp_num_ ,3, sizeof (* r2 ) ,r2_d ,n_warp_num_ ,r2 , n_warp_num_ );//得到梯度2
    }

    //计算M^T*(MX+AP-TPS)
    stat_ = cublasDgemm(handle_,CUBLAS_OP_N,CUBLAS_OP_N,n_allpoint_num,3,n_warp_num_,&beta,m_d,n_allpoint_num,x_d,n_warp_num_,&beta,ap_tps_d,n_allpoint_num);
    stat_ = cublasGetMatrix(n_allpoint_num ,3, sizeof (double) ,ap_tps_d ,n_allpoint_num, tps_error, n_allpoint_num );//tps能量
    
    if(calgrad)
    {
        stat_ = cublasDgemm(handle_,CUBLAS_OP_T,CUBLAS_OP_N,n_warp_num_,3,n_allpoint_num,&ga,m_d,n_allpoint_num,ap_tps_d,n_allpoint_num,&gama,r3_d,n_warp_num_);   
        stat_ = cublasGetMatrix(n_warp_num_ ,3, sizeof (* r3_d ) ,r3_d ,n_warp_num_ ,r3 , n_warp_num_ );//得到梯度3
    }


    //计算梯度大小，直接是下降方向
    if(calgrad)
    {
#pragma omp parallel for  
        for(size_t i =0;i<n_warp_num_*3;i++)
        {
            grad_[i] = -(r1[i]+r2[i]+r3[i]);
        }
    }
    

    //计算能量大小    
    Eigen::VectorXd Vdata;Vdata.resize(n_corr_num_*3);
    Eigen::VectorXd Vsmooth;Vsmooth.resize(n_warp_num_*3);
    Eigen::VectorXd Vtps;Vtps.resize(n_allpoint_num*3);
#pragma omp parallel for    
    for(size_t i =0;i<n_corr_num_*3;i++)
    {
        Vdata(i) = data_error[i];
    }
    data_err = Vdata.squaredNorm()/n_corr_num_;

#pragma omp parallel for  
    for(size_t i =0;i<n_warp_num_*3;i++)
    {
        Vsmooth(i) = smooth_error[i];
    }
    smooth_err = Vsmooth.squaredNorm()/n_warp_num_;

#pragma omp parallel for  
    for(size_t i =0;i<n_allpoint_num*3;i++)
    {
        Vtps(i) = tps_error[i];
    }
    tps_err = Vtps.squaredNorm()/n_allpoint_num;

    //释放指针

    free(r1);
    free(r3);        
    free(r2);
    free(data_error);
    free(smooth_error);
    free(tps_error);
    cudaFree(r1_d);
    cudaFree(r2_d);
    cudaFree(r3_d);
    return (data_err+smooth_err+tps_err);//返回总误差

}

int NRreger::DecentSolver_GPU(Scalar &data_err, Scalar &smooth_err,Scalar &tps_err)
{
    int iter;
    Scalar old_err; 
    Scalar alpha = pars_.learning_rate;
    Scalar *prex;
    prex = (double*)malloc(n_warp_num_*3*sizeof(double));

    
    for (iter = 0; iter < pars_.max_iters; iter++)
    {
        TPSUpdata_GPU(pars_.Tps_R);//更新AP-TPS
        Scalar new_err = sample_gradient_GPU(data_err,smooth_err,tps_err,true);//计算梯度和误差
        std::cout<<"data error:"<<data_err<<"\n";
        std::cout<<"smooth error:"<<smooth_err<<"\n";
        std::cout<<"tps error:"<<tps_err<<"\n";
        //打印下梯度
        // for(int i =0; i<n_warp_num_;i++)
        // {
        //     std::cout<<grad_[i]<<" "<<grad_[i+n_warp_num_]<<" "<<grad_[i+2*n_warp_num_];
        //     std::cout<<std::endl;
        // }


        
        for(size_t i =0;i<n_warp_num_*3;i++)
        {
            prex[i] = x_[i];
            x_[i] = prex[i]+alpha*grad_[i];
        }
        if(iter!=0 && new_err>old_err)
        {
            alpha /= 2;
            for(size_t i =0;i<n_warp_num_*3;i++)
            {
                x_[i] = prex[i]+alpha*grad_[i];
            }
            new_err = sample_gradient_GPU(data_err,smooth_err,tps_err,false);          
        }
        if(new_err-old_err<0.01)
        {
            alpha = alpha*1.2;
        }
        if(new_err-old_err<pars_.stop && new_err < 0.5 )
        {
            break;
        }
        old_err = new_err;

    }
    free(prex);
    return iter;
}

void NRreger::TPS_GPU(PointCloud pc_tps){

    int n = ntps_+3;//系数矩阵的size
    int lda = n;
    
    double *L;//设置系数矩阵
    L = (double*)calloc(n*n,sizeof(double));
    double *Y;//设置系数矩阵
    Y = (double*)calloc(n,sizeof(double));
    // double* X_;//解向量
    X_ = (double*)calloc(n,sizeof(double));
    //给入数据，注意L是一个对称阵
// #pragma omp parallel for 
    for(size_t t= 0; t<ntps_; t++)
    {
        for(size_t j = 0;j<3;j++)
        {
            if(j==0)//第一列是1
            {
                L[(ntps_+j)*n+t] = 1;
                L[t*n+ntps_+j] = 1;
            }
            else
            {
                L[(ntps_+j)*n+t] = pc_tps.points_[t](j-1);
                L[t*n+ntps_+j] = pc_tps.points_[t](j-1);
            }    
        }
        Eigen::Vector2d p1 = {pc_tps.points_[t](0),pc_tps.points_[t](1)};
        for(size_t t2 = 0; t2 < ntps_; t2++)
        {
            if(t!=t2)//轴线上为0,不需要赋值
            {
                Eigen::Vector2d p2 ={pc_tps.points_[t2](0),pc_tps.points_[t2](1)};
                Scalar dis = (p1-p2).norm();
                L[t2*n+t] = std::pow(dis,2)*std::log(dis);
            }
        }
        Y[t] = pc_tps.points_[t](2);
    }


    double *d_L, *d_Y;
    cudaMalloc((void**)&d_L, sizeof(double) * n * n);
    cudaMalloc((void**)&d_Y, sizeof(double) * n );
    cudaMemcpy(d_L, L, sizeof(double) * n * n, cudaMemcpyHostToDevice);
    cudaMemcpy(d_Y, Y, sizeof(double) * n, cudaMemcpyHostToDevice);
    cusolverDnHandle_t cusolverH;
    cusolverDnCreate(&cusolverH);
    int* devIpiv;
    cudaMalloc((void**)&devIpiv, sizeof(int) * n);
    cusolverDnDgetrf(cusolverH, n, n, d_L, lda,NULL,devIpiv,NULL);
    // 解方程
    cusolverDnDgetrs(cusolverH, CUBLAS_OP_N, n, 1, d_L, lda,devIpiv, d_Y, n, NULL);
    cudaMemcpy(X_, d_Y, sizeof(double) * n, cudaMemcpyDeviceToHost);
    cudaFree(d_L);
    cudaFree(d_Y);
    cusolverDnDestroy(cusolverH);

 
}

void NRreger::TPSUpdata_GPU(Scalar R)
{
    
    //计算所有点形变后的位置
    double* deformap_d;
    cudaStat_ = cudaMalloc((void**)&deformap_d,n_allpoint_num*3*sizeof(*deformap));
    double* initialap_d;
    cudaStat_ = cudaMalloc((void**)&initialap_d,n_allpoint_num*3*sizeof(*deformap));
    stat_ = cublasSetMatrix(n_allpoint_num,3,sizeof(*deformap),deformap,n_allpoint_num,deformap_d,n_allpoint_num);
    stat_ = cublasSetMatrix(n_allpoint_num,3,sizeof(*deformap),initialap,n_allpoint_num,initialap_d,n_allpoint_num);
    double al =1.0;
    double beta =1.0;
    stat_ = cublasDgemm(handle_,CUBLAS_OP_N,CUBLAS_OP_N,n_allpoint_num,3,n_warp_num_,&beta,m_d,n_allpoint_num,x_d,n_warp_num_,&beta,initialap_d,n_allpoint_num);//D_temp = Smat_M_*Smat_X_+Smat_AP_;
    stat_ = cublasGetMatrix(n_allpoint_num ,3, sizeof (*deformap_d) ,initialap_d ,n_allpoint_num ,deformap, n_allpoint_num );

    //构造ap-tps
    open3d::geometry::KDTreeFlann* tree = new open3d::geometry::KDTreeFlann(pc_DEM_tar_);

    for(size_t i=0;i<n_allpoint_num;i++)
    {
        Eigen::Vector3d p = {deformap[0*n_allpoint_num+i],deformap[1*n_allpoint_num+i],deformap[2*n_allpoint_num+i]};
        Eigen::Vector2d p2 = {deformap[0*n_allpoint_num+i],deformap[1*n_allpoint_num+i]};
        std::vector<int> index;
        std::vector<Scalar> dis2;
        tree->SearchRadius(p,R,index,dis2);//1表示在多远的范围搜索******非常重要*******
        if(index.size()==0)
        {
            ap_tps_[0*n_allpoint_num+i] = initialap[0*n_allpoint_num+i]-deformap[0*n_allpoint_num+i];
            ap_tps_[1*n_allpoint_num+i] = initialap[1*n_allpoint_num+i]-deformap[1*n_allpoint_num+i];
            ap_tps_[2*n_allpoint_num+i] = initialap[2*n_allpoint_num+i]-deformap[2*n_allpoint_num+i];
        }
        else
        {
            ap_tps_[0*n_allpoint_num+i] = initialap[0*n_allpoint_num+i]-deformap[0*n_allpoint_num+i];
            ap_tps_[1*n_allpoint_num+i] = initialap[1*n_allpoint_num+i]-deformap[1*n_allpoint_num+i];
            // tps插值出z

            double z = X_[ntps_]+X_[ntps_+1]*p(0)+X_[ntps_+2]*p(1);
            for(size_t j = 0;j<ntps_; j++)
            {
                Eigen::Vector2d p1;
                p1<< pc_tps.points_[j](0),pc_tps.points_[j](1);
                double dis = (p1-p2).norm();
                z += std::pow(dis,2)*std::log(dis)*X_[j];
            }
            ap_tps_[2*n_allpoint_num+i] = initialap[2*n_allpoint_num+i]-z;
        }
    }
    stat_ = cublasSetMatrix(n_allpoint_num,3,sizeof(Scalar),ap_tps_,n_allpoint_num, ap_tps_d,n_allpoint_num);


    cudaFree(deformap_d);
    cudaFree(initialap_d);


    
}

open3d::geometry::PointCloud  NRreger::Shower_GPU()
{
    nframe++;
    
    dem_main_->Merge_GPU(deformap,dem_tar_,n_allpoint_num);
    PointCloud pc = dem_main_->Generate_pc();
    Clear();
   
    std::cout<<"当前融合帧数："<< nframe<<"\n";
    std::cout<<"当前点云规模："<<pc.points_.size()<<"\n";
    return pc;


}

